#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# SPDX-License-Identifier: LGPL-2.1-or-later

import argparse
import itertools
import sys
from typing import TextIO

from codegen_utils import c_string_literal

DWARF_INDEX_SECTIONS = (
    ".debug_info",
    ".debug_types",
    ".debug_abbrev",
    ".debug_str",
    ".debug_str_offsets",
    ".debug_line",
    ".debug_line_str",
)

CACHED_SECTIONS = (
    ".debug_addr",
    ".debug_frame",
    ".eh_frame",
    ".orc_unwind_ip",
    ".orc_unwind",
    ".debug_loc",
    ".debug_loclists",
)

UNCACHED_SECTIONS = (
    ".text",
    ".got",
)


def section_enumerator_name(section_name: str) -> str:
    return "DRGN_SCN_" + section_name.lstrip(".").upper()


def gen_elf_sections_h(out_file: TextIO) -> None:
    out_file.write(
        """\
/* Generated by libdrgn/build-aux/gen_elf_sections.py -H. */

#ifndef DRGN_ELF_SECTIONS_H
#define DRGN_ELF_SECTIONS_H

/**
 * Identifiers for important ELF sections so that they can be referenced by
 * index rather than name.
 */
enum drgn_section_index {
"""
    )
    for section_name in DWARF_INDEX_SECTIONS:
        out_file.write(f"\t{section_enumerator_name(section_name)},\n")
    out_file.write(
        """\
	/** Indices less than this are cached when the module is loaded. */
	DRGN_SECTION_INDEX_NUM_PRECACHE,
"""
    )

    for i, section_name in enumerate(CACHED_SECTIONS):
        if i == 0:
            out_file.write(
                f"\t{section_enumerator_name(section_name)} = DRGN_SECTION_INDEX_NUM_PRECACHE,\n"
            )
        else:
            out_file.write(f"\t{section_enumerator_name(section_name)},\n")
    out_file.write(
        """\
	/** Indices less than this may have their data cached. */
	DRGN_SECTION_INDEX_NUM_DATA,
"""
    )

    for i, section_name in enumerate(UNCACHED_SECTIONS):
        if i == 0:
            out_file.write(
                f"\t{section_enumerator_name(section_name)} = DRGN_SECTION_INDEX_NUM_DATA,\n"
            )
        else:
            out_file.write(f"\t{section_enumerator_name(section_name)},\n")
    out_file.write(
        """\
	/** Number of section indices. */
	DRGN_SECTION_INDEX_NUM
};

#endif /* DRGN_ELF_SECTIONS_H */
"""
    )


def gen_drgn_section_name_to_index_inc_strswitch(out_file: TextIO) -> None:
    out_file.write("/* Generated by libdrgn/build-aux/gen_elf_sections.py. */\n")
    out_file.write(
        """\
static enum drgn_section_index drgn_section_name_to_index(const char *name)
{
	@strswitch (name)@
"""
    )
    for section_name in itertools.chain(
        DWARF_INDEX_SECTIONS, CACHED_SECTIONS, UNCACHED_SECTIONS
    ):
        out_file.write(f"\t@case {c_string_literal(section_name)}@\n")
        out_file.write(f"\t\treturn {section_enumerator_name(section_name)};\n")
    out_file.write(
        """\
	@default@
		return DRGN_SECTION_INDEX_NUM;
	@endswitch@
}
"""
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--header", "-H", action="store_true", help="generate header file"
    )
    args = parser.parse_args()

    if args.header:
        gen_elf_sections_h(sys.stdout)
    else:
        gen_drgn_section_name_to_index_inc_strswitch(sys.stdout)


if __name__ == "__main__":
    main()
